#include "BaseDetection.h"
#include "YoloV5Detection.h"

int main()
{
	cv::Mat img = cv::imread("./bus.jpg");
  std::cout << "rows=" << img.rows << " cols=" << img.cols << std::endl;
 
	std::vector<cv::Mat> vec_imgs;
	vec_imgs.emplace_back(img);
	std::vector<std::vector<cv::Rect>> test_result;
  std::string weight_path("./yolov5m.wts");

	//YoloV5Detection yolov5;
	//yolov5.init(weight_path);
	//yolov5.predict(vec_imgs, test_result);
	//yolov5.free();
  BaseDetection* detect = new YoloV5Detection;
 	detect->init(weight_path);
	detect->predict(vec_imgs, test_result, 640, 640);
	detect->free();

	for (int b = 0; b < vec_imgs.size(); b++) {
		auto& res = test_result[b];
		cv::Mat& img = vec_imgs[b];
		for (size_t j = 0; j < res.size(); j++) {
			cv::Rect r = res[j];
			cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
			//cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
		}
		cv::imwrite("./bus_output.jpg", img);
	}
 delete detect;
	return 0;
}